import multiprocessing

import os
import time
from threading import Thread

import numpy as np


class Runner:

    def __init__(self):
        self.tasks = []
        self.isAlive = False

    def append(self, func, args=None, kwargs=None):
        if args is None:
            args = []
        if kwargs is None:
            kwargs = {}
        self.tasks.append((func, args, kwargs))

    def extend(self, func_args_kwargs):
        _func_args_kwargs = []
        for f_a_w in func_args_kwargs:
            assert len(f_a_w) in (1, 2, 3)
            fun = f_a_w[0]
            args = []
            kwargs = {}            
            for a in f_a_w[1:]:
                if isinstance(a, dict):
                    kwargs.update(a)
                else:
                    args.extend(a)
            _func_args_kwargs.append((fun, args, kwargs))

        self.tasks.extend(_func_args_kwargs)

    def pool_apply(self, num_workers=None, ASYNC=True):
        # multiprocessing 进程间无交互
        self.isAlive = True
        if num_workers is None:
            num_workers = len(self.tasks)
        num_workers = min([len(self.tasks), num_workers,
                        multiprocessing.cpu_count()])
        print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} Starting {len(self.tasks)} tasks with {num_workers} processes')
        if not len(self.tasks):
            self.isAlive = False
            return
        pool = multiprocessing.Pool(processes=num_workers)
        for i, (func, args, kwargs) in enumerate(self.tasks, 1):
            # print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} Process {i} starting {func.__name__}')
            if ASYNC:
                pool.apply_async(func, args, kwargs)
            else:
                pool.apply(func, args, kwargs)

        pool.close()
        pool.join()
        self.isAlive = False
        print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} {len(self.tasks)} tasks finished')

    def thread_start(self, num_workers=None, ASYNC=True):
        # multithreading 线程间无交互
        # assert 一个cpu核能开启多少线程？
        self.isAlive = True
        if num_workers is None:
            num_workers = len(self.tasks)
        num_workers = min([len(self.tasks), num_workers,
                        multiprocessing.cpu_count() * 4])
        print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} Starting {len(self.tasks)} tasks with {num_workers} threads')
        if not len(self.tasks):
            self.isAlive = False
            return

        def start(i, tasks):
            thread_list = []

            for i, (func, args, kwargs) in enumerate(tasks, i):
                # print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} Thread {i} starting {func.__name__}')
                thread = Thread(target=func, args=args, kwargs=kwargs)
                thread.start()
                if not ASYNC:
                    thread.join()
                thread_list.append(thread)
            while True:
                time.sleep(1.0)
                isAlive = False
                for t in thread_list:
                    if t.is_alive():
                        isAlive = True
                if not isAlive:
                    break
            self.isAlive = False
                
        for i in range(0, len(self.tasks), num_workers):
            tasks = self.tasks[i: i + num_workers]
            start(i, tasks)

        print(f'{time.strftime("%Y-%m-%d %H:%M:%S")} {len(self.tasks)} tasks finished')

    def run(self, mode='process', num_workers=None, ASYNC=True):
        assert mode in ('process', 'thread')

        exec_ = self.pool_apply if mode == 'process' else self.thread_start
        exec_(num_workers, ASYNC)


if __name__ == '__main__':
    from tqdm import tqdm

    def fun(n):
        hexcodes = []
        tbar = tqdm(n)
        for i in tbar:
            hexcode = 1 # hexdigest(
            #     (np.random.randn(128, 128, 2)*255).astype(np.uint8))
            tbar.desc = f'{i} {hexcode}'
            hexcodes.append(hexcode)

        print(len(hexcodes), len(set(hexcodes)))

    def fun(a, b, c):
        print('model', a)
        print('args', b, c)

    runner = Runner()
    # runner.extend([[fun, [n] * n] for n in [100, 200, 300, 400, 500]])
    runner.extend([(fun, (a, b, c), {}) for a, b, c in (('model1', 'args1', 'args2'), ('model2', 'args1', 'args2'), ('model3', 'args1', 'args2'))])
    # runner.run('process', ASYNC=False)
    runner.run('thread', num_workers=1, ASYNC=False)